# Load Lasso data to get optimal lambda
for (mod in c("lasso","rf","nn","gbm")) {
  load(paste(modeldir,"/",
             country,
             "_",mod,"_",
             "count",
             "_",
             rhs.group,fileext,
             ".RData",
             sep = ""))
}

#####################################
######### B- GET ESTIMATES ##########
#####################################
#####################################

# Set holder for each year's results/model
ebma.results <- list()
# Loop over years
i = 0
for (year in start.year:end.year) { 
  i = i + 1
  print(year)
  # Generate dependent variable
  dta.past <- dta[dta[,t.var]< year,]
  N <- length(dta.past[,dv])
  ###################################
  ### 1 # Find Optimal Parameters ###
  ###################################
  
  # Do repeated cross validation to find best weights and DT
  cv_results <- foreach (cv=1:cv.runs, 
                         .combine='rbind',
                         .errorhandling = 'remove',
                         .packages = c("glmnet",
                                       "randomForest",
                                       "gbm",
                                       "nnet",
                                       "pROC")) %dopar% {   # Loop over CV runs
                                         set.seed(cv)
                                         shuffle <- sample(1:N,
                                                           N,
                                                           replace=FALSE)
                                         sampleSplit = c(0,round((N/5*c(1:5))))
    # Set matrix to store predictions
    pred <- matrix(NA,nrow=N,ncol=5)
    colnames(pred) <- c("lasso.prediction",
                        "rf.prediction",
                        "gbm.prediction",
                        "nn.prediction",
                        "ebma.prediction")
    for (f in 1:5) {     # Loop over folds       
      # Fit lasso
      lasso.fit <-  glmnet(x = as.matrix(dta.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                                 rhs]),
                                     y = as.matrix(dta.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                                 dv]),
                                     alpha = lasso_count_params$alpha,
                                     family = lasso_count_params$family,
                                     lambda = (seq(991,1,-10)^2)*lasso.results[[i]]$lambda)
      opt.lam <- 100
      while (is.na(lasso.fit$lambda[opt.lam])) {
        opt.lam = opt.lam - 1
      }
      # Save predictions for fold f
      pred[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
           "lasso.prediction"] <- predict(lasso.fit,
                                                          newx = as.matrix(dta.past[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                                                         rhs]),
                                                          type = "response")[,opt.lam]
      # Fit Random Forest
      rf.fit <- randomForest(x = as.matrix(dta.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                         rhs]),
                             y = as.matrix(dta.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                         dv]),
                             ntree = rf_count_params$ntree, 
                             nodesize = rf_count_params$nodesize, 
                             type = rf_count_params$type,
                             do.trace = rf_count_params$do.trace
      )
      # Save predictions for fold f
      pred[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],"rf.prediction"] <- predict(rf.fit,
                                                       newdata = as.matrix(dta.past[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                                                         rhs]))
      # Fit AdaBoost
      gbm.fit <- gbm.fit(x = as.matrix(dta.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                     rhs]),
                         y = as.matrix(dta.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                               dv]),
                         n.trees = gbm.results[[i]]$trees,
                         distribution = gbm_count_params$distribution,
                         interaction.depth = gbm_count_params$interaction.depth,
                         n.minobsinnode = gbm_count_params$n.minobsinnode,
                         verbose = gbm_count_params$verbose,
                         shrinkage = gbm_count_params$shrinkage,
                         train.fraction = gbm_count_params$train.fraction,
                         bag.fraction = gbm_count_params$bag.fraction)
      # Save predictions for fold f
      pred[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],"gbm.prediction"] <- predict(gbm.fit,
                                                        newdata = as.matrix(dta.past[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                                                          rhs]),
                                                        type = "response",
                                                        n.trees = gbm.results[[i]]$trees)
      
      # Fit Neural Network
      # Remove constant variables from predictors
      train.RHS <- dta.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                 rhs]
      test.RHS <- dta.past[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                rhs]
      # Remove variables that are constant
      not.constant <- apply(train.RHS,2,sd)!=0
      train.RHS <- train.RHS[,not.constant]
      test.RHS <- test.RHS[,not.constant]
      # Find PC weights
      pcs.fit <- prcomp(train.RHS,center = TRUE, scale = TRUE)
      train.pcs <- predict(pcs.fit,
                           train.RHS)[,1:min(nn_count_params$pcNum,
                                             ncol(train.RHS))]
      test.pcs <- predict(pcs.fit,
                          test.RHS)[,1:min(nn_count_params$pcNum,
                                           ncol(train.RHS))]
      # Fit Network
      nn.fit <- nnet(x = train.pcs,
                     y = as.matrix(dta.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                 dv]),
                     size = nn.results[[i]]$size,
                     decay = nn_count_params$decay,
                     trace = nn_count_params$trace,
                     linout = nn_count_params$linout
      )
      
      # Save predictions for fold f
      pred[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],"nn.prediction"] <- predict(nn.fit,
                                                       newdata = test.pcs)
      
    }
    pred[pred<0] <- 0
    # Get weights for each model
    l.contrib <- (((pred[,c("lasso.prediction",
                                   "rf.prediction",
                                   "gbm.prediction",
                                   "nn.prediction")]^dta.past[,dv])
                   * exp(-pred[,c("lasso.prediction",
                                         "rf.prediction",
                                         "gbm.prediction",
                                         "nn.prediction")]))
                  /factorial(dta.past[,dv]))
    weight.denom <- colSums(l.contrib)
    weight <- weight.denom/sum(weight.denom)
    pred[,"ebma.prediction"] <- as.matrix(pred[,c("lasso.prediction",
                                                            "rf.prediction",
                                                            "gbm.prediction",
                                                            "nn.prediction")]) %*% as.matrix(weight)
    cv_results <- c(cv,weight)
    names(cv_results) <- c("CV Run",
                           "lasso.weight",
                           "rf.weight",
                           "gbm.weight",
                           "nn.weight") 
    return(cv_results)
  } # End cv runs
  hours <- floor((proc.time()[3]-start.time)/3600)
  mins <- floor((proc.time()[3]-start.time)/60) - hours*60
  secs <- floor((proc.time()[3]-start.time)) - hours*3600 - mins*60
  print(paste(hours,
              "h",
              mins,
              "m",
              secs,
              "s elapsed",
              sep = ""))
  
  ######################################
  ### 2 # FIT OPTIMAL MODEL ############
  ######################################
  
  # A1 - Fit Model
  ###############
  
  ebma.results[[i]] <- list()
  ebma.results[[i]]$optim.wts <- colMeans(cv_results[,c("lasso.weight",
                                                    "rf.weight",
                                                    "gbm.weight",
                                                    "nn.weight")])
  # B1 Fit out of sample
  ########################
  
  ebma.results[[i]]$outsample.component.prob <- cbind("lasso" = lasso.results[[i]]$fit.oos,
                                                      "rf" = rf.results[[i]]$fit.oos,
                                                      "gbm" = gbm.results[[i]]$fit.oos,
                                                      "nn" = nn.results[[i]]$fit.oos)
  ebma.results[[i]]$fit.oos <- ebma.results[[i]]$outsample.component.prob %*% as.matrix(ebma.results[[i]]$optim.wts)
  ebma.results[[i]]$actual.oos <- dta[dta[,t.var]== year,dv]
}
save(ebma.results,
     file=paste(modeldir,"/",
                country,
                "_ebma_",
                "count",
                "_",
                rhs.group,fileext,
                ".RData",
                sep = ""))